{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练词向量\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## 利用`gensim.models.word2vec`训练词向量\n",
    "- 原始语料为[中文维基](https://dumps.wikimedia.org/zhwiki/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "0. 原始语料为 `xml` 格式，需要提取出正文，使用 `WikiExtractor` 包\n",
    "    0. 命令行提取正文：`python WikiExtractor.py -b 500M -o wiki zhwiki-20190720-pages-articles-multistream.xml.bz2`\n",
    "    0. 获得的文件中，正文被包含在 `<doc></doc>` 标签内\n",
    "0. 或者 `gensim.corpora.WikiCorpus` 直接处理 `xml.bz2` 文件\n",
    "0. 由上两步，获得的文本先经过预处理，**每一行一句话，单词间用空格隔开**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python WikiExtractor.py -b 500M -o datasets/wiki datasets/zhwiki-20190720.xml.bz2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "def preprocess_zhwiki_v1():\n",
    "    # 提取文本信息，分句、分词、繁体转简体，然后将单词用空格连接\n",
    "    regex = re.compile(\"(^<doc.*>$)|(^</doc>$)\")\n",
    "    sent_spliter = re.compile(\"。|！|？\")\n",
    "\n",
    "    input_file = open(input_file_path, 'r', encoding='utf-8')\n",
    "    output_file = open(output_file_path, 'w+', encoding='utf-8')\n",
    "\n",
    "    line = input_file.readline()\n",
    "    while line:\n",
    "        if line.strip() and not regex.match(line):\n",
    "            sentences = sent_spliter.split(line)\n",
    "            for s in sentences:\n",
    "                s = zhconv.convert(s, 'zh-cn')\n",
    "                words = jieba.cut(s.strip('\\n'))\n",
    "                sent = ' '.join(words)\n",
    "                output_file.write(sent + '\\n')\n",
    "        line = input_file.readline()\n",
    "\n",
    "    input_file.close()\n",
    "    output_file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gensim.corpora import WikiCorpus\n",
    "\n",
    "\n",
    "def preprocess_zhwiki_v2():\n",
    "    # 提取文本信息，分句、分词、繁体转简体，然后将单词用空格连接\n",
    "    # WikiCorpus 会将标点符号都被删除\n",
    "    space = ' '\n",
    "    i = 0\n",
    "    l = []\n",
    "\n",
    "    output_file = open(output_file_path, 'w+', encoding='utf-8')\n",
    "\n",
    "    wiki = WikiCorpus(input_file_path, lemmatize=False, dictionary={})\n",
    "    for text in wiki.get_texts():\n",
    "        for temp_sentence in text:\n",
    "            temp_sentence = zhconv.convert(s, 'zh-cn')\n",
    "            seg_list = list(jieba.cut(temp_sentence))\n",
    "            for temp_term in seg_list:\n",
    "                l.append(temp_term)\n",
    "        output_file.write(space.join(l) + '\\n')\n",
    "        l = []\n",
    "        i = i + 1\n",
    "\n",
    "        if (i % 200 == 0):\n",
    "            print('Saved ' + str(i) + ' articles')\n",
    "    f.close()\n",
    "\n",
    "\n",
    "input_file_path = r'datasets/wiki/AA/wiki_00'\n",
    "output_file_path = r'datasets/wiki/AA/wiki_corpus'\n",
    "preprocess_zhwiki_v2()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 利用上一步生成的处理后的满足 `LineSentence` 格式的文本，创建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gensim.models import word2vec\n",
    "\n",
    "corpus_path = output_file_path\n",
    "model_path = r\"models/wiki_corpus.model\"\n",
    "\n",
    "\n",
    "def build_model(corpus_path):\n",
    "    wiki_news = word2vec.LineSentence(corpus_path)\n",
    "    model = word2vec.Word2Vec(\n",
    "        wiki_news,\n",
    "        sg=0,  # 模型类型 CBOW\n",
    "        size=50,  # 词向量维度     \n",
    "        window=5,  # 窗口尺寸\n",
    "        min_count=5, # 忽略词频少于 5 的单词\n",
    "        workers=9)\n",
    "    model.save(model_path)\n",
    "    return model\n",
    "\n",
    "\n",
    "build_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 验证训练得到的模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gensim.models import word2vec\n",
    "model_path = r\"models/zhwiki.50d.word2vec\"\n",
    "model = word2vec.Word2Vec.load(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.wv.most_similar('数学')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.wv.most_similar('哲学')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.wv.most_similar(positive=['女人', '国王'], negative=['男人'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "two_corpus = [\"腾讯\",\"阿里巴巴\"]\n",
    "res = model.wv.similarity(two_corpus[0],two_corpus[1])\n",
    "print(\"similarity:%.4f\"%res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 将词向量降维后进行可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import TruncatedSVD\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "%matplotlib inline\n",
    "\n",
    "word_vectors = model.wv\n",
    "\n",
    "\n",
    "def get_model_matrix(word_vectors, required_words):\n",
    "    import random\n",
    "    words = list(word_vectors.vocab.keys())\n",
    "    random.shuffle(words)\n",
    "    words = words[:10000]\n",
    "    print(\"Putting %i words into word2Ind and matrix M...\" % len(words))\n",
    "    word2Ind = {}\n",
    "    M = []\n",
    "    curInd = 0\n",
    "    for w in words:\n",
    "        try:\n",
    "            M.append(word_vectors.word_vec(w))\n",
    "            word2Ind[w] = curInd\n",
    "            curInd += 1\n",
    "        except KeyError:\n",
    "            continue\n",
    "    for w in required_words:\n",
    "        try:\n",
    "            M.append(word_vectors.word_vec(w))\n",
    "            word2Ind[w] = curInd\n",
    "            curInd += 1\n",
    "        except KeyError:\n",
    "            continue\n",
    "    M = np.stack(M)\n",
    "    print(\"Done.\")\n",
    "    return M, word2Ind\n",
    "\n",
    "\n",
    "words = [\n",
    "    '数学', '算术', '公理', '积分', '统计', '善恶', '哲学', '伦理', '中国政府', '美国国会', '武侠小说',\n",
    "    '风靡', '海内外', '受欢迎', '通俗小说', '中华人民共和国', '文化大革命', '反思', '伤痕', '一批', '白话文',\n",
    "    '诗人', '古诗', '欢迎', '中华民国', '撤退', '台湾', '区别', '思潮', '过渡时期', '通称', '文献', '兴趣',\n",
    "    '钻研', '语言学', '神秘主义', '更加', '经典', '历史学', '文学', '学术界', '享有', '前所未有', '趋势',\n",
    "    '受到', '人文主义者', '巨量', '规则', '机器人', '精准', '身躯', '脑', '视频', '确保', '高质量', '适中',\n",
    "    '价格', '软件设计', '构成', '互补', '并行', '系统分析', '程序设计', '支持', '高级', '课程', '训练',\n",
    "    '工业', '技能', '羧酸', '柠檬酸', '高效率', '肽键', '细胞骨架', '细胞周期', '氯仿', '甘油', '变型',\n",
    "    '鞘', '类固醇', '醛', '酮', '糖原', '单糖', '半乳糖', '葡萄糖', '糖苷键', '含氮', '杂环', '嘌呤',\n",
    "    '辅酶', '底物', '化学能', '磷酸化', '哈康', '延斯', '挪威海', '捕鲸', '挪威政府', '成人礼', '巴伦支海',\n",
    "    '哥德堡', '区域规划', '润州', '邳州市', '东海县', '丹阳市', '武进区', '临河', '嘈杂', '霰弹枪', '讲席',\n",
    "    '一滴', '调换', '香港金融管理局', '美圆', '金管局', '毫', '大额', '铜币', '一圆', '镍币', '爆竹',\n",
    "    '管理科', '中区', '收兑', '财政司'\n",
    "]\n",
    "\n",
    "M, word2Ind = get_model_matrix(word_vectors, words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 利用 svd 算法进行降维\n",
    "def reduce_to_k_dim(M, k=2):\n",
    "    n_iters = 10\n",
    "    print(\"Running Truncated SVD over %i words...\" % (M.shape[0]))\n",
    "    svd = TruncatedSVD(n_components=k, n_iter=n_iters)\n",
    "    M_reduced = svd.fit_transform(M)\n",
    "    print(\"Done.\")\n",
    "    return M_reduced\n",
    "\n",
    "M_reduced = reduce_to_k_dim(M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_embeddings(M_reduced, word2Ind, words):\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(24,24))\n",
    "    for word in words:\n",
    "        index = word2Ind[word]\n",
    "        x, y = M_reduced[index]\n",
    "        plt.scatter(x, y, marker='o', color='red')\n",
    "        plt.text(x, y, word, fontsize=9)\n",
    "        \n",
    "plot_embeddings(M_reduced, word2Ind, words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 利用 TSNE 算法进行降维\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "def tsne_plot(M, word2Ind, words):\n",
    "\n",
    "    tsne_model = TSNE(perplexity=40, n_components=2, init='pca', n_iter=2500, random_state=23)\n",
    "    M_reduced = tsne_model.fit_transform(M)\n",
    "\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(32, 32))\n",
    "    for word in words:\n",
    "        index = word2Ind[word]\n",
    "        x, y = M_reduced[index]\n",
    "        ax.scatter(x, y, marker='o', color='red')\n",
    "        ax.text(x, y, word, fontsize=9)\n",
    "        \n",
    "tsne_plot(M, word2Ind, words)\n",
    "\n",
    "# TSNE 降维效果比 SVD 要好，但效率更低"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 关键词提取，从 `wv.most_similar()` 出发获取给定单词的相关单词\n",
    "    - 词向量 `wv.most_similar()` 获得的为出现在相似上下文中的同类词，并不是通常语义含义上的相似词"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "\n",
    "def get_related_words(initial_words, model):\n",
    "    unseen = [initial_words]\n",
    "    seen = defaultdict(int)\n",
    "\n",
    "    max_size = 500\n",
    "\n",
    "    while unseen and len(seen) < max_size:\n",
    "        if len(seen) % 50 == 0:\n",
    "            print('search length: {}'.format(len(seen)))\n",
    "\n",
    "        node = unseen.pop(0)\n",
    "        new_expanding = [w for w, _ in model.most_similar(node, topn=20)]\n",
    "        unseen += new_expanding\n",
    "\n",
    "        seen[node] += 1\n",
    "    return seen\n",
    "\n",
    "\n",
    "actions = get_related_words(\"说\", word_vectors)\n",
    "actions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `wordcloud` 实现词云"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 培根散文集的词云\n",
    "data_path = r'datasets/Bacon Francis - Essays.txt'\n",
    "\n",
    "import os\n",
    "\n",
    "from os import path\n",
    "from wordcloud import WordCloud\n",
    "\n",
    "text = open(data_path).read()\n",
    "\n",
    "wordcloud = WordCloud().generate(text)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plt.imshow(wordcloud, interpolation='bilinear')\n",
    "plt.axis(\"off\")\n",
    "\n",
    "wordcloud = WordCloud(max_font_size=40).generate(text)\n",
    "plt.figure()\n",
    "plt.imshow(wordcloud, interpolation=\"bilinear\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `tensorflow.nn.sampled_softmax_loss` 训练词向量"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 利用上述预处理后的文件\n",
    "- 需要根据语料创建词汇表\n",
    "- 需要将句子分词后的词语列表，转变成 中心词-上下文词 组成的词对"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T10:04:53.852955Z",
     "start_time": "2020-04-08T10:04:52.641593Z"
    }
   },
   "outputs": [],
   "source": [
    "import zipfile\n",
    "import os\n",
    "from collections import Counter\n",
    "import random\n",
    "import bz2\n",
    "\n",
    "import random\n",
    "import math\n",
    "from six.moves import xrange\n",
    "\n",
    "import zipfile\n",
    "import tensorflow as tf\n",
    "\n",
    "import collections\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "tf.test.is_gpu_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T10:06:37.547294Z",
     "start_time": "2020-04-08T10:06:35.964139Z"
    }
   },
   "outputs": [],
   "source": [
    "# 将原始文件转化成单词列表\n",
    "# 大段的文本直接利用空格拆分\n",
    "def read_data(file_path):\n",
    "    with zipfile.ZipFile(file_path) as f:\n",
    "        words = tf.compat.as_text(f.read(f.namelist()[0])).split()\n",
    "    return words\n",
    "\n",
    "\n",
    "filename = r'datasets/text8.zip'\n",
    "vocabulary = read_data(filename)\n",
    "\n",
    "# 文本中总单词量\n",
    "print('Data size', len(vocabulary))\n",
    "\n",
    "# 词汇表大小\n",
    "len(set(vocabulary))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T10:08:23.788429Z",
     "start_time": "2020-04-08T10:08:20.825628Z"
    }
   },
   "outputs": [],
   "source": [
    "# 创建词典，替换稀有词\n",
    "vocabulary_size = 10000\n",
    "\n",
    "\n",
    "def build_dataset(words, n_words):\n",
    "    # 单词-词频\n",
    "    count = [['UNK', -1]]\n",
    "    count.extend(collections.Counter(words).most_common(n_words - 1))\n",
    "\n",
    "    # 词典：单词-索引\n",
    "    dictionary = dict()\n",
    "    for word, _ in count:\n",
    "        dictionary[word] = len(dictionary)\n",
    "\n",
    "    # 文本向量化\n",
    "    data = list()\n",
    "    unk_count = 0\n",
    "    for word in words:\n",
    "        index = dictionary.get(word, 0)\n",
    "        if index == 0:\n",
    "            unk_count += 1\n",
    "        data.append(index)\n",
    "    count[0][1] = unk_count\n",
    "\n",
    "    # 索引-单词，词典\n",
    "    reversed_dictionary = dict(zip(dictionary.values(), dictionary.keys()))\n",
    "    return data, count, dictionary, reversed_dictionary\n",
    "\n",
    "\n",
    "data, count, dictionary, reverse_dictionary = build_dataset(\n",
    "    vocabulary, vocabulary_size)\n",
    "del vocabulary\n",
    "\n",
    "print('Most common words (+UNK)', count[:5])\n",
    "print('Sample data', data[:10], [reverse_dictionary[i] for i in data[:10]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 生成批处理数据\n",
    "    - 将文本`the quick brown fox jumped over the lazy dog`转换为数组 $[123,6,53,24,5,12,89,8,11]$\n",
    "    - $skip\\_window = 2$，对于中心词 `53` 则上下文-中心词对$ [(123，53),(6，53),(24，53),(5,53]$\n",
    "    - $num\\_skips = 3 $，从每个窗口词对中选择 3 个，如 $[(123,53),(24,53),(5,53)]$\n",
    "    - 遍历所有中心词，将词对转化成解包，如 batch 为中心词数组 $[53,53,53]$，labels 为上下文词数组 $[123, 24, 5]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 生成用于 skip-gram 模型的训练数据\n",
    "data_index = 0\n",
    "\n",
    "\n",
    "def generate_batch(batch_size, num_skips, skip_window):\n",
    "    global data_index\n",
    "    assert batch_size % num_skips == 0\n",
    "    assert num_skips <= 2 * skip_window\n",
    "    batch = np.ndarray(shape=(batch_size), dtype=np.int32)\n",
    "    labels = np.ndarray(shape=(batch_size, 1), dtype=np.int32)\n",
    "    span = 2 * skip_window + 1\n",
    "    buffer = collections.deque(maxlen=span)\n",
    "    if data_index + span > len(data):\n",
    "        data_index = 0\n",
    "    buffer.extend(data[data_index:data_index + span])\n",
    "    data_index += span\n",
    "    for i in range(batch_size // num_skips):\n",
    "        context_words = [w for w in range(span) if w != skip_window]\n",
    "        words_to_use = random.sample(context_words, num_skips)\n",
    "        for j, context_word in enumerate(words_to_use):\n",
    "            batch[i * num_skips + j] = buffer[skip_window]\n",
    "            labels[i * num_skips + j] = buffer[context_word]\n",
    "        if data_index == len(data):\n",
    "            buffer.extend(data[0:span])\n",
    "            data_index = span\n",
    "        else:\n",
    "            buffer.append(data[data_index])\n",
    "            data_index += 1\n",
    "    data_index = (data_index + len(data) - span) % len(data)\n",
    "    return batch, labels\n",
    "\n",
    "\n",
    "batch, labels = generate_batch(batch_size=8, num_skips=2, skip_window=2)\n",
    "for i in range(8):\n",
    "    print(batch[i], reverse_dictionary[batch[i]], '->', labels[i, 0],\n",
    "          reverse_dictionary[labels[i, 0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[reverse_dictionary[i] for i in batch]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 模型参数\n",
    "\n",
    "batch_size = 128\n",
    "embedding_size = 64  # 词向量的维度\n",
    "skip_window = 1  # 窗口大小\n",
    "num_skips = 2  # 每个窗口的词对中选择多少对\n",
    "num_sampled = 64  # 负采样时负样本数量\n",
    "\n",
    "valid_size = 16  # Random set of words to evaluate similarity on.\n",
    "valid_window = 100  # Only pick dev samples in the head of the distribution.\n",
    "valid_examples = np.random.choice(valid_window, valid_size, replace=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建模型，tf.nn.sampled_softmax_loss\n",
    "initializer_softmax = tf.keras.initializers.GlorotUniform()\n",
    "# Variables:\n",
    "embeddings_weight = tf.Variable(\n",
    "    tf.random.uniform([vocabulary_size, embedding_size], -1.0, 1.0))\n",
    "softmax_weight = tf.Variable(\n",
    "    initializer_softmax([vocabulary_size, embedding_size]))\n",
    "softmax_bias = tf.Variable(initializer_softmax([vocabulary_size]))\n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "\n",
    "num_step = 100001\n",
    "for step in range(num_step):\n",
    "    batch_inputs, batch_labels = generate_batch(batch_size, num_skips,\n",
    "                                                skip_window)\n",
    "    batch_inputs = tf.cast(batch_inputs, tf.int32)\n",
    "    batch_labels = tf.cast(batch_labels, tf.int32)\n",
    "\n",
    "    with tf.GradientTape() as tape:\n",
    "        embed = tf.nn.embedding_lookup(embeddings_weight, batch_inputs)\n",
    "        loss = tf.reduce_mean(\n",
    "            tf.nn.sampled_softmax_loss(weights=softmax_weight,\n",
    "                                       biases=softmax_bias,\n",
    "                                       inputs=embed,\n",
    "                                       labels=batch_labels,\n",
    "                                       num_sampled=num_sampled,\n",
    "                                       num_classes=vocabulary_size))\n",
    "    variables = [embeddings_weight, softmax_weight, softmax_bias]\n",
    "    gradients = tape.gradient(loss, variables)\n",
    "    optimizer.apply_gradients(zip(gradients, variables))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings_weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:47:22.596016Z",
     "start_time": "2020-04-08T12:47:22.590897Z"
    }
   },
   "source": [
    "## `tensorflow.nn.nce_loss` 训练词向量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建模型\n",
    "batch_size = 64\n",
    "vocabulary_size = len(word2ind)\n",
    "embedding_dimension = 5\n",
    "negative_samples = 8\n",
    "LOG_DIR = 'logs/word2vec'\n",
    "\n",
    "embeddings = tf.Variable(\n",
    "    tf.random.uniform([vocabulary_size, embedding_dimension], -1.0, 1.0))\n",
    "nce_weights = tf.Variable(\n",
    "    tf.random.truncated_normal([vocabulary_size, embedding_dimension],\n",
    "                               stddev=1.0 / math.sqrt(embedding_dimension)))\n",
    "nce_biases = tf.Variable(tf.zeros([vocabulary_size]))\n",
    "\n",
    "variables = [embeddings, nce_weights, nce_weights]\n",
    "\n",
    "\n",
    "def loss_fn(embed, labels):\n",
    "    return tf.reduce_mean(\n",
    "        tf.nn.nce_loss(weights=nce_weights,\n",
    "                       biases=nce_biases,\n",
    "                       inputs=embed,\n",
    "                       labels=labels,\n",
    "                       num_sampled=negative_samples,\n",
    "                       num_classes=vocabulary_size))\n",
    "\n",
    "\n",
    "optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)\n",
    "\n",
    "for i in range(1000):\n",
    "    x_batch, y_batch = generate_batch(batch_size)\n",
    "    x_batch = tf.constant(x_batch, dtype=tf.int32)\n",
    "    y_batch = tf.constant(y_batch, dtype=tf.int32)\n",
    "    with tf.GradientTape() as tape:\n",
    "        embed = tf.nn.embedding_lookup(embeddings, x_batch)\n",
    "        loss = loss_fn(embed, y_batch)\n",
    "    grads = tape.gradient(loss, variables)\n",
    "    optimizer.apply_gradients(zip(grads, variables))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `tensorflow`负采样训练词向量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:05.981831Z",
     "start_time": "2020-04-08T12:31:05.165551Z"
    }
   },
   "outputs": [],
   "source": [
    "import collections\n",
    "import math\n",
    "import random\n",
    "import sys\n",
    "import time\n",
    "import os\n",
    "import numpy as np\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T10:13:32.234593Z",
     "start_time": "2020-04-08T10:13:32.226666Z"
    }
   },
   "source": [
    "### 处理数据集\n",
    "华尔街日报的文章，每一行一个句子，句子中每个词由空格隔开"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:06.088490Z",
     "start_time": "2020-04-08T12:31:05.991310Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "42068\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['aer',\n",
       " 'banknote',\n",
       " 'berlitz',\n",
       " 'calloway',\n",
       " 'centrust',\n",
       " 'cluett',\n",
       " 'fromstein',\n",
       " 'gitano',\n",
       " 'guterman',\n",
       " 'hydro-quebec',\n",
       " 'ipo',\n",
       " 'kia',\n",
       " 'memotec',\n",
       " 'mlx',\n",
       " 'nahb',\n",
       " 'punts',\n",
       " 'rake',\n",
       " 'regatta',\n",
       " 'rubens',\n",
       " 'sim',\n",
       " 'snack-food',\n",
       " 'ssangyong',\n",
       " 'swapo',\n",
       " 'wachter']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert 'ptb.train.txt' in os.listdir('datasets/ptb')\n",
    "\n",
    "# 句子列表，每个句子为单词组成的列表\n",
    "with open('datasets/ptb/ptb.train.txt', 'r') as f:\n",
    "    lines = f.readlines()\n",
    "    raw_dataset = [st.split() for st in lines]\n",
    "\n",
    "print(len(raw_dataset))\n",
    "raw_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:09.987423Z",
     "start_time": "2020-04-08T12:31:09.981617Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# tokens: 24 ['aer', 'banknote', 'berlitz', 'calloway', 'centrust']\n",
      "# tokens: 15 ['pierre', '<unk>', 'N', 'years', 'old']\n",
      "# tokens: 11 ['mr.', '<unk>', 'is', 'chairman', 'of']\n"
     ]
    }
   ],
   "source": [
    "for st in raw_dataset[:3]:\n",
    "    print('# tokens:', len(st), st[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 建立词语索引"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:10.979899Z",
     "start_time": "2020-04-08T12:31:10.900348Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(9858, 25)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 统计单词出现的频次，删除词频<5的单词\n",
    "counter = collections.Counter([tk for st in raw_dataset for tk in st])\n",
    "counter = dict(filter(lambda x: x[1] >= 5, counter.items()))\n",
    "len(counter), counter['learn']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:11.715725Z",
     "start_time": "2020-04-08T12:31:11.617748Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "887100"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 创建单词与索引映射\n",
    "idx2token = [tk for tk, _ in counter.items()]\n",
    "token2idx = {tk: idx for idx, tk in enumerate(idx2token)}\n",
    "\n",
    "# 将数据集转化为索引列表\n",
    "dataset = [[token2idx[tk] for tk in st if tk in token2idx]\n",
    "           for st in raw_dataset]\n",
    "\n",
    "# 语料总单词树\n",
    "num_tokens = sum([len(st) for st in dataset])\n",
    "num_tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 二次采样\n",
    "文本数据中一般会出现一些高频词，如英文中的“the”“a”和“in”，在一个背景窗口中，一个词（如“chip”）和较低频词（如“microprocessor”）同时出现比和较高频词（如“the”）同时出现对训练词嵌入模型更有益。因此，训练词嵌入模型时可以对词进行二次采样。 具体来说，数据集中每个被索引词$w_i$将有一定概率被丢弃，该丢弃概率为\n",
    "$$ P(w_i) = \\max\\left(1 - \\sqrt{\\frac{t}{f(w_i)}}, 0\\right)$$\n",
    "其中 $f(w_i)$ 是数据集中词$w_i$的个数与总词数之比，常数$t$是一个超参数（实验中设为$10^{-4}$）。可见，只有当$f(w_i) > t$时，我们才有可能在二次采样中丢弃词$w_i$，并且越高频的词被丢弃的概率越大。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:14.367013Z",
     "start_time": "2020-04-08T12:31:13.970089Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'# tokens: 375402'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 是否丢弃单词，该单词在语料中出现的总次数越高，越可能被丢弃\n",
    "def discard(idx):\n",
    "    return random.uniform(\n",
    "        0, 1) < 1 - math.sqrt(1e-4 / counter[idx2token[idx]] * num_tokens)\n",
    "\n",
    "# 重采样后数据集\n",
    "subsampled_dataset = [[tk for tk in st if not discard(tk)] for st in dataset]\n",
    "\n",
    "# 语料总单词数大大减少\n",
    "'# tokens: %d' % sum([len(st) for st in subsampled_dataset])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:16.978815Z",
     "start_time": "2020-04-08T12:31:16.940305Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# the: before=50770, after=2114\n",
      "# join: before=45, after=45\n"
     ]
    }
   ],
   "source": [
    "# 二次采样后，单词被保留的次数\n",
    "def compare_counts(token):\n",
    "    return '# %s: before=%d, after=%d' % (\n",
    "        token, sum([st.count(token2idx[token]) for st in dataset]),\n",
    "        sum([st.count(token2idx[token]) for st in subsampled_dataset]))\n",
    "\n",
    "\n",
    "print(compare_counts('the'))\n",
    "# 高频词the被保留的次数从 50770 降低到 2153\n",
    "print(compare_counts('join'))\n",
    "# 低频词join基本都保留"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 提取中心词和背景词\n",
    "每次在整数1和max_window_size（最大背景窗口）之间随机均匀采样一个整数作为背景窗口大小。每个中心词，对应一个背景词列表"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:18.073395Z",
     "start_time": "2020-04-08T12:31:18.067059Z"
    }
   },
   "outputs": [],
   "source": [
    "# 每个中心词对应一个背景词列表\n",
    "def get_centers_and_contexts(dataset, max_window_size):\n",
    "    centers, contexts = [], []\n",
    "    for st in dataset:\n",
    "        if len(st) < 2:  # 每个句子至少要有2个词才可能组成一对“中心词-背景词”\n",
    "            continue\n",
    "        centers += st  # 每句话所有单词都是中心词\n",
    "        for i in range(len(st)):\n",
    "            window_size = random.randint(1, max_window_size)\n",
    "            indices = list(\n",
    "                range(max(0, i - window_size), min(len(st),\n",
    "                                                   i + window_size + 1)))\n",
    "            indices.remove(i)\n",
    "            contexts.append([st[idx] for idx in indices])\n",
    "    return centers, contexts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:18.892452Z",
     "start_time": "2020-04-08T12:31:18.885003Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset [[0, 1, 2, 3, 4], [8, 9]]\n",
      "center 0 has contexts [1, 2]\n",
      "center 1 has contexts [0, 2, 3]\n",
      "center 2 has contexts [1, 3]\n",
      "center 3 has contexts [1, 2, 4]\n",
      "center 4 has contexts [3]\n",
      "center 8 has contexts [9]\n",
      "center 9 has contexts [8]\n"
     ]
    }
   ],
   "source": [
    "tiny_dataset = [list(range(5)), list(range(8, 10))]\n",
    "print('dataset', tiny_dataset)\n",
    "for center, context in zip(*get_centers_and_contexts(tiny_dataset, 2)):\n",
    "    print('center', center, 'has contexts', context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:22.551101Z",
     "start_time": "2020-04-08T12:31:21.754140Z"
    }
   },
   "outputs": [],
   "source": [
    "all_centers, all_contexts = get_centers_and_contexts(subsampled_dataset, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 负采样\n",
    "- 常规的$softmax$会输出整个词汇表的概率分布，计算量巨大\n",
    "- 使用负采样来进行近似训练：对于一对中心词和背景词，我们随机采样$K$个噪声词（实验中设$K=5$）。根据word2vec论文的建议，噪声词采样概率$P(w)$设为$w$词频与总词频之比的0.75次方。\n",
    "- 采样之后，一个中性词对应一个背景词列表，还对应一个噪声词列表，噪声词列表的长度是背景词列表长度的$K$倍，此时$sofmax$输出为背景词+噪声词列表的概率分布"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:29.268617Z",
     "start_time": "2020-04-08T12:31:23.332296Z"
    }
   },
   "outputs": [],
   "source": [
    "# 以词频的0.75次方作为权重进行采样\n",
    "sampling_weights = [counter[w]**0.75 for w in idx2token]\n",
    "\n",
    "\n",
    "def get_negatives(all_contexts, sampling_weights, K):\n",
    "    all_negatives, neg_candidates, i = [], [], 0\n",
    "    population = list(range(len(sampling_weights)))\n",
    "    for contexts in all_contexts:\n",
    "        negatives = []\n",
    "        while len(negatives) < len(contexts) * K:\n",
    "            if i == len(neg_candidates):\n",
    "                # 循环开始时，根据每个词的权重（sampling_weights）随机生成k个词的索引作为候选噪声词。\n",
    "                # 为了高效计算，可以将k设得稍大一点\n",
    "                i, neg_candidates = 0, random.choices(population,\n",
    "                                                      sampling_weights,\n",
    "                                                      k=int(1e5))\n",
    "            # 依次将候选噪声词添加到噪声词中\n",
    "            neg, i = neg_candidates[i], i + 1\n",
    "            # 噪声词不能是背景词\n",
    "            if neg not in set(contexts):\n",
    "                negatives.append(neg)\n",
    "        all_negatives.append(negatives)\n",
    "    return all_negatives\n",
    "\n",
    "\n",
    "all_negatives = get_negatives(all_contexts, sampling_weights, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:31:31.361295Z",
     "start_time": "2020-04-08T12:31:31.327317Z"
    }
   },
   "outputs": [],
   "source": [
    "# 每个中心词对应的噪声词列表的长度，是背景词列表长度的𝐾倍\n",
    "assert all([\n",
    "    len(neg) == len(window) * 5\n",
    "    for window, neg in zip(all_contexts, all_negatives)\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据管道\n",
    "- 每个单词，对应一个**相同长度**的上下文词列表：上下文由背景词+噪声词+填充组成；\n",
    "- masks 表明上下文中单词是否是填充；\n",
    "- labels 表明上下文中的每个单词是不是背景词"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:01.305519Z",
     "start_time": "2020-04-08T12:34:01.298323Z"
    }
   },
   "outputs": [],
   "source": [
    "def preprocess(centers, contexts, negatives):\n",
    "    max_len = max(\n",
    "        len(cont) + len(neg) for cont, neg in zip(contexts, negatives))\n",
    "    targets, masks, labels = [], [], []\n",
    "    for cont, neg in zip(contexts, negatives):\n",
    "        cur_len = len(cont) + len(neg)\n",
    "\n",
    "        # 用 0 填充成相同的长度\n",
    "        targets.append(cont + neg + [0] * (max_len - cur_len))\n",
    "\n",
    "        # 区别单词和填充\n",
    "        masks.append([1] * cur_len + [0] * (max_len - cur_len))\n",
    "\n",
    "        # 背景词才是标签，噪声词和填充都不是\n",
    "        labels.append([1] * len(cont) + [0] * (max_len - len(cont)))\n",
    "    centers = np.array(centers, dtype=np.float32).reshape(-1, 1)\n",
    "    targets = np.array(targets, dtype=np.float32)\n",
    "    masks = np.array(masks, dtype=np.float32)\n",
    "    labels = np.array(labels, dtype=np.float32)\n",
    "    return (centers, targets, masks, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:08.207501Z",
     "start_time": "2020-04-08T12:34:03.254115Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.82 s, sys: 136 ms, total: 4.95 s\n",
      "Wall time: 4.95 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "dataset = tf.data.Dataset.from_tensor_slices(preprocess(\n",
    "    all_centers, all_contexts, all_negatives))\n",
    "\n",
    "batch_size = 512\n",
    "dataset = dataset.shuffle(len(all_centers)).batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:10.968984Z",
     "start_time": "2020-04-08T12:34:10.133893Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "centers shape: (512, 1)\n",
      "contexts_negatives shape: (512, 60)\n",
      "masks shape: (512, 60)\n",
      "labels shape: (512, 60)\n"
     ]
    }
   ],
   "source": [
    "for batch in dataset.take(1):\n",
    "    for name, data in zip(['centers', 'contexts_negatives', 'masks', 'labels'],\n",
    "                          batch):\n",
    "        print(name, 'shape:', data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:14.189613Z",
     "start_time": "2020-04-08T12:34:13.381274Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "centers shape: <dtype: 'float32'>\n",
      "contexts_negatives shape: <dtype: 'float32'>\n",
      "masks shape: <dtype: 'float32'>\n",
      "labels shape: <dtype: 'float32'>\n"
     ]
    }
   ],
   "source": [
    "for batch in dataset.take(1):\n",
    "    for name, data in zip(['centers', 'contexts_negatives', 'masks', 'labels'],\n",
    "                          batch):\n",
    "        print(name, 'shape:', data.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `skip-gram`模型\n",
    "利用中心词预测(背景词+噪声词)的概率分布"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:18.009051Z",
     "start_time": "2020-04-08T12:34:17.995818Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[ 0.00248704,  0.04588217,  0.03903292,  0.02340596],\n",
       "        [ 0.01726497, -0.04898253,  0.03221582,  0.02792713],\n",
       "        [-0.03564789,  0.04589545, -0.04923713, -0.00122647],\n",
       "        [ 0.04300623,  0.03403026,  0.03677107,  0.02419565],\n",
       "        [ 0.01740289, -0.03229471,  0.00015046,  0.00863054],\n",
       "        [ 0.04152843,  0.04772225,  0.02492184,  0.01654983]],\n",
       "       dtype=float32)]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 嵌入层\n",
    "embed = tf.keras.layers.Embedding(input_dim=6, output_dim=4)\n",
    "embed.build(input_shape=(1, 6))\n",
    "embed.get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:22.807498Z",
     "start_time": "2020-04-08T12:34:22.800852Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=77, shape=(2, 3, 4), dtype=float32, numpy=\n",
       "array([[[ 0.01726497, -0.04898253,  0.03221582,  0.02792713],\n",
       "        [-0.03564789,  0.04589545, -0.04923713, -0.00122647],\n",
       "        [ 0.04300623,  0.03403026,  0.03677107,  0.02419565]],\n",
       "\n",
       "       [[ 0.01740289, -0.03229471,  0.00015046,  0.00863054],\n",
       "        [ 0.04152843,  0.04772225,  0.02492184,  0.01654983],\n",
       "        [ 0.00248704,  0.04588217,  0.03903292,  0.02340596]]],\n",
       "      dtype=float32)>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = tf.constant([[1, 2, 3], [4, 5, 0]], dtype=tf.int32)\n",
    "embed(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:24.736197Z",
     "start_time": "2020-04-08T12:34:24.733754Z"
    }
   },
   "outputs": [],
   "source": [
    "# 跳字模型，批量数据，输出权重向量\n",
    "@tf.function\n",
    "def skip_gram(centers, contexts, embed_v, embed_u):\n",
    "    v = embed_v(centers)\n",
    "    u = embed_u(contexts)\n",
    "    pred = tf.matmul(v, tf.transpose(u, perm=[0, 2, 1]))\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T11:28:39.800967Z",
     "start_time": "2020-04-08T11:28:39.795864Z"
    }
   },
   "source": [
    "### 训练模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 损失函数\n",
    "负采样后的数据中引入了`mask`，填充的数据在训练时应该被遮蔽掉，对应的损失不应该计入模型损失函数。损失函数可以使用二元交叉熵损失函数，下面定义`SigmoidBinaryCrossEntropyLoss`:`x = logits`, `z = labels`\n",
    "\n",
    "$$loss = z\\times(-log(sigmoid(x))) + (1 - z)\\times(-log(1 - sigmoid(x)))$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:32.947876Z",
     "start_time": "2020-04-08T12:34:32.941824Z"
    }
   },
   "outputs": [],
   "source": [
    "class SigmoidBinaryCrossEntropy(tf.keras.losses.Loss):\n",
    "    def __init__(self):\n",
    "        super(SigmoidBinaryCrossEntropy, self).__init__()\n",
    "\n",
    "    def __call__(self, labels, logits, mask=None):\n",
    "        res = tf.nn.sigmoid_cross_entropy_with_logits(labels, logits) * mask\n",
    "        return tf.reduce_mean(res, axis=1)\n",
    "\n",
    "\n",
    "loss = SigmoidBinaryCrossEntropy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:38.122627Z",
     "start_time": "2020-04-08T12:34:37.753109Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=98, shape=(2,), dtype=float32, numpy=array([0.8739896, 1.2099689], dtype=float32)>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred = tf.constant([[1.5, 0.3, -1, 2], [1.1, -0.6, 2.2, 0.4]],\n",
    "                   dtype=tf.float32)\n",
    "\n",
    "# 标签变量label中的1和0分别代表背景词和噪声词\n",
    "label = tf.constant([[1, 0, 0, 0], [1, 1, 0, 0]], dtype=tf.float32)\n",
    "\n",
    "# 掩码变量\n",
    "mask = tf.constant([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=tf.float32)\n",
    "\n",
    "loss(label, pred, mask) * mask.shape[1] / tf.reduce_sum(mask, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:34:51.391785Z",
     "start_time": "2020-04-08T12:34:51.370811Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.layers.embeddings.Embedding at 0x7fd5165667d0>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 第一个嵌入层为需要的词向量，第二个嵌入层丢弃\n",
    "embed_size = 100\n",
    "net = tf.keras.Sequential([\n",
    "    tf.keras.layers.Embedding(input_dim=len(idx2token),\n",
    "                              output_dim=embed_size,\n",
    "                              name='word2vec'),\n",
    "    tf.keras.layers.Embedding(input_dim=len(idx2token),\n",
    "                              output_dim=embed_size,\n",
    "                              name='output')\n",
    "])\n",
    "net.get_layer(name='word2vec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:35:01.622175Z",
     "start_time": "2020-04-08T12:35:01.614331Z"
    }
   },
   "outputs": [],
   "source": [
    "# 训练过程\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)\n",
    "\n",
    "def train(num_epochs):\n",
    "    for epoch in range(num_epochs):\n",
    "        start, cur_loss, n = time.time(), 0.0, 0\n",
    "        for center, target, mask, label in dataset:\n",
    "            with tf.GradientTape(persistent=True) as tape:\n",
    "                pred = skip_gram(center, target, net.get_layer(index=0),\n",
    "                                 net.get_layer(index=1))\n",
    "                # 使用掩码变量mask来避免填充项对损失函数计算的影响\n",
    "                l = (loss(label, tf.reshape(pred, label.shape), mask) *\n",
    "                     mask.shape[1] / tf.reduce_sum(mask, axis=1))\n",
    "                l = tf.reduce_mean(l)  # 一个batch的平均loss\n",
    "\n",
    "            grads = tape.gradient(l, net.variables)\n",
    "            optimizer.apply_gradients(zip(grads, net.variables))\n",
    "            cur_loss += l.numpy().item()\n",
    "            n += 1\n",
    "        print('epoch %d, loss %.2f, time %.2fs' %\n",
    "              (epoch + 1, cur_loss / n, time.time() - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:38:35.462984Z",
     "start_time": "2020-04-08T12:35:04.478511Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1, loss 0.45, time 21.91s\n",
      "epoch 2, loss 0.39, time 21.23s\n",
      "epoch 3, loss 0.35, time 21.08s\n",
      "epoch 4, loss 0.32, time 20.95s\n",
      "epoch 5, loss 0.31, time 20.93s\n",
      "epoch 6, loss 0.30, time 20.90s\n",
      "epoch 7, loss 0.30, time 20.95s\n",
      "epoch 8, loss 0.29, time 20.91s\n",
      "epoch 9, loss 0.29, time 21.21s\n",
      "epoch 10, loss 0.28, time 20.91s\n"
     ]
    }
   ],
   "source": [
    "train(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 验证模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-08T12:39:41.592862Z",
     "start_time": "2020-04-08T12:39:41.579251Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosine sim=0.534: microprocessors\n",
      "cosine sim=0.505: dell\n",
      "cosine sim=0.498: micro\n",
      "cosine sim=0.457: folk\n",
      "cosine sim=0.447: intel\n"
     ]
    }
   ],
   "source": [
    "def get_similar_tokens(query_token, k, embed):\n",
    "    W = embed.get_weights()\n",
    "    W = tf.convert_to_tensor(W[0])\n",
    "    x = W[token2idx[query_token]]\n",
    "    x = tf.reshape(x, shape=[-1, 1])\n",
    "    # 添加的1e-9是为了数值稳定性\n",
    "    cos = tf.reshape(tf.matmul(W, x), shape=[\n",
    "        -1\n",
    "    ]) / tf.sqrt(tf.reduce_sum(W * W, axis=1) * tf.reduce_sum(x * x) + 1e-9)\n",
    "    _, topk = tf.math.top_k(cos, k=k + 1)\n",
    "    topk = topk.numpy().tolist()\n",
    "    for i in topk[1:]:  # 除去输入词\n",
    "        print('cosine sim=%.3f: %s' % (cos[i], (idx2token[i])))\n",
    "\n",
    "\n",
    "get_similar_tokens('chip', 5, net.get_layer(index=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "165px"
   },
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "position": {
    "height": "879.85px",
    "left": "433px",
    "right": "20px",
    "top": "120px",
    "width": "358.5px"
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
